from PySide6.QtWidgets import (QDockWidget, QTreeWidget, QTreeWidgetItem, 
                             QMenu, QFileDialog,QDialog)
from PySide6.QtCore import Qt, Signal
from PySide6.QtGui import QIcon
from ui.data_import_dialog import DataImportDialog
import os
import io
import re
import pandas as pd
import chardet
from core.data_model import DataModel

class ProjectTreeDock(QDockWidget):
    visualization_requested = Signal(str)  # 新增信号
    open_table_requested = Signal(str)  # 信号：井名称
    dock_area = Qt.LeftDockWidgetArea

    def __init__(self, parent=None):
        super().__init__("项目结构", parent)
        self.setFeatures(QDockWidget.DockWidgetClosable | 
                        QDockWidget.DockWidgetMovable | 
                        QDockWidget.DockWidgetFloatable)
        
        self.tree = QTreeWidget()
        self.tree.setHeaderHidden(True)
        self.tree.setDragEnabled(True)
        self.tree.setAcceptDrops(True)
        self.tree.setDropIndicatorShown(True)
        self.tree.setDragDropMode(QTreeWidget.InternalMove)
        
        self.setWidget(self.tree)
        self.data_model = DataModel()
        self._init_tree()
        
        # 连接信号
        self.tree.itemDoubleClicked.connect(self._on_item_double_clicked)
        self.tree.setContextMenuPolicy(Qt.CustomContextMenu)
        self.tree.customContextMenuRequested.connect(self._show_context_menu)

    def _init_tree(self):
        """初始化项目树"""
        # 创建根节点
        root = QTreeWidgetItem(self.tree, ["井数据"])
        root.setIcon(0, QIcon(":/icons/project.png"))
        
        # 从数据库加载井列表
        wells = self.data_model.get_well_list()
        for well in wells:
            item = QTreeWidgetItem(root, [well['well_name']])
            item.setIcon(0, QIcon(":/icons/well.png"))
            # 存储井的详细信息
            item.setData(0, Qt.UserRole, well)
        
        self.tree.expandAll()

    def _on_item_double_clicked(self, item, column):
        """处理双击事件"""
        if item.parent() and item.parent().text(0) == "井数据":
            well_name = item.text(0)
            self.visualization_requested.emit(well_name)

    def _show_context_menu(self, pos):
        """显示右键菜单"""
        item = self.tree.itemAt(pos)
        if not item:
            return
            
        menu = QMenu(self)
        
        if item.parent() and item.parent().text(0) == "井数据":
            open_action = menu.addAction("打开数据")
            open_action.triggered.connect(lambda: self.open_table_requested.emit(item.text(0)))
            
            menu.addSeparator()
            
            export_action = menu.addAction("导出数据")
            export_action.triggered.connect(lambda: self._export_data(item))
        else:
            refresh_action = menu.addAction("刷新")
            refresh_action.triggered.connect(self._refresh_tree)

        menu.addAction("导入数据").triggered.connect(self._import_data)

        menu.exec_(self.tree.mapToGlobal(pos))

    def _export_data(self, item):
        """导出数据"""
        well_name = item.text(0)
        data = self.data_model.get_data(well_name)
        if data is not None:
            file_path, _ = QFileDialog.getSaveFileName(
                self, "导出数据", f"{well_name}.xlsx", "Excel Files (*.xlsx)"
            )
            if file_path:
                data.to_excel(file_path, index=False)

    def _refresh_tree(self):
        """刷新项目树"""
        self.tree.clear()
        self._init_tree() 

    def _add_folder(self):
        # Implementation of adding a folder
        pass

    def _add_file(self):
        # Implementation of adding a file
        pass

    # def _import_data(self):
    #     # Implementation of importing data
    #     pass
    # def _import_data(self):
    #     file_dialog = QFileDialog()
    #     file_path, _ = file_dialog.getOpenFileName(self, "选择文件", "", "文本文件 (*.txt *.csv)")
    #     if not file_path:
    #         return
    #
    #     try:
    #         # 提取井名
    #         file_name = os.path.basename(file_path)
    #         well_name = self.extract_well_name(file_name)
    #
    #         # 检测编码并读取数据
    #         with open(file_path, 'rb') as f:
    #             raw_data = f.read()
    #             result = chardet.detect(raw_data)
    #             encoding = result['encoding'] or 'gbk'
    #             if encoding in ['ISO-8859-1', 'ascii']:
    #                 encoding = 'gbk'
    #
    #         df = self.try_three_formats(file_path, encoding)
    #
    #         if df.empty:
    #             print("警告：读取的数据为空")
    #             return
    #
    #         # 更新树结构
    #         root = self.tree.invisibleRootItem()
    #         for i in range(root.childCount()):
    #             item = root.child(i)
    #             if item.text(0) == "井数据":
    #                 # 创建新井节点
    #                 well_item = QTreeWidgetItem(item, [well_name])
    #                 well_item.setIcon(0, QIcon(":/icons/well.png"))
    #                 curve_collection = QTreeWidgetItem(well_item, ["曲线集合"])
    #                 for col in df.columns:
    #                     child_item = QTreeWidgetItem(curve_collection, [col])
    #                     child_item.setCheckState(0, Qt.Unchecked)
    #                 self.tree.expandAll()
    #                 break
    #
    #     except Exception as e:
    #         print(f"读取文件时出错: {e}")

    def _import_data(self):
        file_dialog = QFileDialog()
        file_path, _ = file_dialog.getOpenFileName(self, "选择文件", "", "文本文件 (*.txt *.csv)")
        if not file_path:
            return

        try:
            # 提取井名
            file_name = os.path.basename(file_path)
            well_name = self.extract_well_name(file_name)

            # 检测编码并读取数据
            with open(file_path, 'rb') as f:
                raw_data = f.read()
                result = chardet.detect(raw_data)
                encoding = result['encoding'] or 'gbk'
                if encoding in ['ISO-8859-1', 'ascii']:
                    encoding = 'gbk'

            df = self.try_three_formats(file_path, encoding)

            if df.empty:
                print("警告：读取的数据为空")
                return

            # 弹出导入配置窗口
            dialog = DataImportDialog(file_path, df, parent=self)
            if dialog.exec_() == QDialog.Accepted:
                processed_df = dialog.get_processed_data()
                # 更新数据模型
                self.data_model.add_data(well_name, processed_df)
                # 更新项目树
                self._update_tree_with_fields(well_name, processed_df.columns.tolist())

        except Exception as e:
            print(f"读取文件时出错: {e}")

    def extract_well_name(self, file_name):
        pattern = r'(\d+-\d+[A-Za-z]+)'
        match = re.search(pattern, file_name)
        if match:
            return match.group(1) + '井'
        index = file_name.find('井')
        if index != -1:
            return file_name[:index + 1]
        return file_name.split('.')[0]

    def try_three_formats(self, file_path, encoding):
        # 格式1: 井深 钻时 全烃（简单表格）
        try:
            df = pd.read_csv(file_path, sep=r'\s+', header=0, encoding=encoding)
            if not df.empty:
                return df
        except:
            pass

        # 格式2: 井号 深度 岩性简述（简单表格）
        try:
            df = pd.read_csv(file_path, sep=r'\s+', header=0, encoding=encoding)
            if not df.empty:
                return df
        except:
            pass

        # 格式3: 包含 ~ASCII DATA SECTION（复杂格式）
        with open(file_path, 'r', encoding=encoding) as f:
            lines = f.readlines()
            data_start = next((i for i, line in enumerate(lines) if line.strip().startswith('~ASCII DATA SECTION')),
                              None)
            if data_start is None:
                return pd.DataFrame()

            # 提取表头（假设表头在 DATA SECTION 前2行）
            header_line = lines[data_start - 2].strip().lstrip('#').strip()
            headers = header_line.split() if header_line else []

            # 提取数据行
            data_lines = [line for line in lines[data_start + 1:] if line.strip()]
            if not data_lines:
                return pd.DataFrame()

            # 解析数据
            data_str = '\n'.join(data_lines)
            df = pd.read_csv(io.StringIO(data_str), sep=r'\s+', header=None, names=headers)
            return df

    def _on_visualization_requested(self, well_name):
        data = self.data_model.get_data(well_name)
        if data is not None:
            self.visualization.update_data(well_name, data)
            self.property_dock.update_data(well_name, data) 